#!/usr/bin/env python
#-*- coding: utf-8 -*-

# Author: Edward Roualdes
#         University of Kentucky

from __future__ import print_function
from Bio import SeqIO
import getopt
import sys
import os
import errno
import re
import shutil
from collections import defaultdict

def parseCommandLine(argv):
    u"""extract specified command line options

Parse command line options and offer help / suggestions if arguments given are incorrectly. 

Input:
    argv: some command line options.

Output:
    outputfile: string. path to the specified config file.
    """
    
    outputfile = ""
    try:
        options, remainder = getopt.getopt(sys.argv[1:], 
                                           'hc:v', 
                                           ['help', 'configfile=', 'version'])
    except getopt.GetoptError as err:
        print(str(err))
        print("Try instead '$ python NextAllele.py -c path/to/CONFIGFILE'")
        sys.exit(2)
    for opt, arg in options:
        if opt in ('-h', '--help'):
            print("Usage")
            print("\tpython NextAllele.py [FLAGS] [CONFIGFILE]\n")
            print("Flags")
            print("{:<20} {:<}".format("\t-h, --help", "print this help message and exit"))
            print("{:<20} {:<}".format("\t-v, --version", "print version number and exit"))
            print("{:<20} {:<}".format("\t-c, --configfile", "specify CONFIGFILE: -c path/to/CONFIGFILE"))
            sys.exit()
        elif opt in ('-v', '--version'):
            print("NextAllele v0.1")
            sys.exit()
        elif opt in ('-c', '--configfile'):
            outputfile = arg
    print("\nCONFIGFILE =", os.path.abspath(outputfile))
    return outputfile


def getNames(FILE):
    u"""get read names from fasta file *FILE*

Parse fasta file *FILE* for read names via SeqIO.parse()

Input:
    FILE: string. path to some file that is readable by SeqIO.parse().

Output:
    Names: a set of read names.
    """

    Names = set()
    file = open(FILE, "rU")
    for seq in SeqIO.parse(file, "fasta"): 
        Names.add(seq.name)
    file.close()

    return Names


def mkdir(path, overwrite, opt = None):
    u"""make directory as specified by the path *path*

Make directory with the user prompted option to *overwrite* existing directory tree, or exit program to specify a new directory. If the user refuses to overwrite the directory tree, *opt* is printed along with the declarative to rerun NextAllele.

Input:
    path: string. path to some (possibly extant) directory.
    overwrite: boolean. True => overwrite extant directory tree.
    opt: a string (=None) to provide a specific error message in case you choose not to overwrite the directory *path*.

Output:
    void.
    """

    try:
        os.makedirs(path)       # just try to make it
    except OSError as exc:      # catch exception
        if exc.errno == errno.EEXIST:
            if overwrite:
                # if exists, choice
                inp = raw_input(path + " already exists, do you want to overwrite it and all of its contents? ")
                res = re.match("[Yy][Ee][Ss]|[Yy]", inp)
                if res:             # delete directory and re-make it
                    shutil.rmtree(path)
                    os.makedirs(path)
                else: # give option kill NextAllele and specify a new OUTPUTDIRECTORY
                    print(opt + "and rerun NextAllele.")
                    sys.exit()
            else:
                pass
        else:
            raise
    return


def writeSeqs(sortedBarcodes, outDir):
    u"""write sorted sequences to files
    
Write the sorted sequences *sortedBarcodes* to files, one folder per individual one file per locus, within directory *outDir*.

Input:
    sortedBarcodes: a defaultdict of {barcode_name_1: [dnaRead_1_1, ..., dnaRead_1_m], barcode_name_c: [dnaRead_c_1, ..., dnaRead_c_l]} # though the lists of reads are possibly tuples ...?
    outDir: string. path to a directory.

Output:
    void.
    """

    for inds, reads in sortedBarcodes.iteritems():
        for read in reads:      # for each read
            individual = inds
            locus = read.annotations["Locus"]

            # make dir each individual w/o overwriting
            mkdir(os.path.join(outDir, individual), False)

            # one file per locus
            seqFile = open(os.path.join(outDir, individual, locus+'.fna'), 'a')

            # write sequences
            SeqIO.write(read, seqFile, "fasta")
            seqFile.close()
    print("Wrote alignments to file...")


# save fasta files in directory to dictionaries, yield one at a time
def files2dict(DIRECTORY, grep, pat = None):
    u"""walk directory *DIRECTORY* and iteratively yield a dictionary for each grep matched fasta file

Walk directory *DIRECTORY* and iteratively yield a dictionary (k = sequence name, v = sequence) for each string pattern *grep* matched fasta file within *DIRECTORY* and any subdirectories. Individual and locus names found via directory's path structure: directory => (indiv1 => loci1), (indiv1 => loci2), (indiv1 => lociK), ... , (indivN => lociM).

Input:
    DIRECTORY: string. path to some directory containing fasta files.
    grep: string pattern. a grep string pattern.
    pat: a string that is searched for in case any primary alignments have been modified.

Output:
    a triplet of a {sequence name: sequence} dict and the corresponding locus and individual names    
    """
    grep = re.compile(grep)
    if not pat:
        # use 'modified' if no alternative keyword is supplied
        # why is this here and not elsewhere?
        pat = 'modified'

    for root, dirs, files in os.walk(DIRECTORY):
        # find and use only files matching grep
        for fname in filter(grep.match, files):

            f = os.path.join(root, fname)  # use absolute path names
            indiv = os.path.basename(root) # save indiv name
            locus = os.path.splitext(os.path.splitext(fname)[0])[0] # save locus name
            data = {}
            # look for "modified" version of each aligned locus file 
            for r, ds, fs in os.walk(root):
                for fsname in filter(lambda x: locus in x, fs):
                    if pat in fsname: # if it exists
                        f = os.path.join(root, fsname) # read from it

            alignedFile = open(f, "rU")
            data = SeqIO.to_dict(SeqIO.parse(alignedFile, "fasta"))
            alignedFile.close()
            yield data, indiv, locus

def unique(seq):
    u"""return list of unique elemets in *seq*

A list of unique elements in *seq* with the order of the elements maintained.

Inputs:
    seq: a list. 

Output:
    result: a list of unique elements from *seq* in the same order.
"""

    seen = {}
    result = []
    for item in seq:
        if item in seen: continue
        seen[item] = 1
        result.append(item)
    return result

def read_matrix(data):
    u"""turn the dictionary *data* into a list of lists (= a row-major matrix of reads).

Input:
    data = {read_1.desc: seq_1, ..., read_n.desc: seq_n}

Output:
    mat = [seq_1, ..., seq_n]     # where seq_k  = ['a', 'c', ..., 'g', 't']
    names = [read_1.desc, ..., read_n.desc] # read_k.desc = 'description of read'

Assumptions:
    m = len(seq_k) for all k = 1, ..., n
"""
    mat = []
    names = []
    for k, v in data.iteritems():
        mat.append(list(v))
        names.append(k)

    # all rows must have same length
    row_lens = [len(r) for r in mat]
    max_row_len = max(row_lens)
    min_row_len = min(row_lens)
    assert max_row_len == min_row_len, "Not all rows have same length!"

    return (mat, names)

def let2num(x):
    u"""change letters to numbers for appropriate STRUCTURE/PHASE input."""
    # check for capital letters, since I don't know
    # if they will make it this far
    if x == 'a' or x == 'A':
        return str(1)
    elif x == 'c' or x == 'C':
        return str(2)
    elif x == 'g' or x == 'G':
        return str(3)
    elif x == 't' or x == 'T':
        return str(4)
    elif x == '-' or x == 'N' or x == 'n': 
        return str(-9)      
    else:
        print('an unexpected letter!', x)
        sys.exit()


if __name__ == "__main__":
    parseCommandLine(sys.argv[1:])
